from PIL import Image
import os
from torch.utils.data import Dataset


class MyDataReader(Dataset):

    def __init__(self, root_path, is_train, transforms):
        self.root_path = root_path
        if is_train:
            self.root_path = os.path.join(self.root_path, "train")
            self.images = os.listdir(self.root_path)
        else:
            self.root_path = os.path.join(self.root_path, "test")
            self.images = os.listdir(self.root_path)
        self.transforms = transforms

    def __getitem__(self, idx):
        image_name = self.images[idx]
        image_item_path = os.path.join(self.root_path, image_name)

        # 打开图片
        image = Image.open(image_item_path)

        # 标签处理
        label = int(image_name.split('_')[2].split('.')[0])  # 转换为整数标签

        # 图片处理
        image = self.transforms(image)

        return image, label

    def __len__(self):
        return len(self.images)
